import tensorflow as tf
import matplotlib.pyplot as plt
import numpy as np
from sklearn.utils import shuffle

# 加载MNIST数据集，这里先把所有数据合并到一起
mnist = tf.keras.datasets.mnist
(X_all, y_all), _ = mnist.load_data()

# 数据归一化，明确转换数据类型为float32，确保后续操作不出问题
X_all = X_all.astype(np.float32) / 255.0
# 增加通道维度，使用更明确的reshape方式
X_all = X_all.reshape(X_all.shape[0], X_all.shape[1], X_all.shape[2], 1)

print("X_all shape:", X_all.shape)
y_all = y_all.reshape(-1, 1).astype(np.float32)  # 转换标签数据类型与图像数据类型一致，同时调整维度
print("y_all shape:", y_all.shape)

# 使用tile函数正确复制维度，使y_all维度与X_all匹配，便于后续合并
y_all_expanded = np.tile(y_all, (1, X_all.shape[1], X_all.shape[2], 1))
print("y_all_expanded shape:", y_all_expanded.shape)

# 将图像数据和标签数据合并到一起，方便后续一起打乱
combined_data = np.concatenate([X_all, y_all_expanded], axis=1)

# 随机打乱数据
shuffled_data = shuffle(combined_data)

# 重新划分训练集和测试集，比如按照80%训练、20%测试的比例（可自行调整比例）
split_index = int(len(shuffled_data) * 0.8)
train_data = shuffled_data[:split_index]
test_data = shuffled_data[split_index:]

# 分离出图像和标签，仔细核对维度计算
X_train = train_data[:, :28 * 28 * 1].reshape(-1, 28, 28, 1)
y_train = train_data[:, 28 * 28 * 1:].reshape(-1, 1)
X_test = test_data[:, :28 * 28 * 1].reshape(-1, 28, 28, 1)
y_test = test_data[:, 28 * 28 * 1:].reshape(-1, 1)

# 打印数据维度，用于检查是否符合预期
print("X_train shape:", X_train.shape)
print("y_train shape:", y_train.shape)
print("X_test shape:", X_test.shape)
print("y_test shape:", y_test.shape)

model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)),
    tf.keras.layers.MaxPooling2D((2, 2)),
    tf.keras.layers.Conv2D(64, (3, 3), activation='relu'),
    tf.keras.layers.MaxPooling2D((2, 2)),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])

model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
model.fit(X_train, y_train, epochs=5, validation_data=(X_test, y_test))

# 模型评估
loss, accuracy = model.evaluate(X_test, y_test)
print(f'Accuracy: {accuracy}')

# 预测
predictions = model.predict(X_test)

# 绘制前10个测试样本及其预测结果
for i in range(10):
    plt.subplot(2, 5, i + 1)
    plt.imshow(X_test[i].reshape(28, 28), cmap='gray')
    plt.title(f'Pred: {np.argmax(predictions[i])}, True: {y_test[i][0]}')
    plt.axis('off')
plt.show()